import openai
from openai import OpenAI  # 新式调用
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from neo4j import GraphDatabase
import json
import time
'''
此模块提供了一个函数，用于从文本中提取问题和回答，
并将这些信息导入到Neo4j数据库中。
'''
# ======================
# 配置区（重点修改这里）
# ======================
NEO4J_URI = "bolt://localhost:7687"
NEO4J_USER = "neo4j"
NEO4J_PASSWORD = "12345678"

# ✅ OpenAI 兼容接口配置
OPENAI_BASE_URL = "https://api.siliconflow.cn/v1"  # 替换为你的接口地址
OPENAI_API_KEY = "sk-xpczcmiojslifyfbfzkrkctbgbilljxoekpevjhcdnrcdzjb"                         # 替换为你的 key
OPENAI_MODEL = "Qwen/Qwen3-235B-A22B-Instruct-2507"                               # 模型名，看接口支持什么

# 向量相似度 top-k
TOP_K = 10

# ======================
# 初始化客户端
# ======================
# Neo4j
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))

# OpenAI 兼容接口客户端
client = OpenAI(
    api_key=OPENAI_API_KEY,
    base_url=OPENAI_BASE_URL
)


# ======================
# 从 Neo4j 获取数据（不变）
# ======================
def get_all_answers():
    with driver.session() as session:
        result = session.run("""
            MATCH (a:Answer)
            RETURN a.id AS id, a.text AS text, a.embedding AS embedding
        """)
        return [record.data() for record in result]


def get_all_tables():
    with driver.session() as session:
        result = session.run("""
            MATCH (t:Table)
            RETURN t.name AS name, t.schema AS schema, 
                   t.comment AS comment, t.embedding AS embedding
        """)
        return [record.data() for record in result]


# ======================
# 向量相似度筛选（不变）
# ======================
def find_topk_candidate_tables(answer_embedding, tables, top_k=TOP_K):
    answer_vec = np.array(answer_embedding).reshape(1, -1)
    table_vecs = np.array([t['embedding'] for t in tables])
    sims = cosine_similarity(answer_vec, table_vecs)[0]
    topk_indices = np.argsort(sims)[-top_k:][::-1]
    return [tables[i] for i in topk_indices]


# ======================
# 调用兼容接口（关键修改）
# ======================
def get_relevant_table_names(answer_text, candidate_tables):
    table_info = [
        {
            "name": t['name'],
            "schema": t['schema'],
            "comment": t.get('comment', '无描述')
        }
        for t in candidate_tables
    ]

    prompt = f"""
你是一个数据库语义分析助手。请根据回答内容，判断它是否需要从以下数据库表中查询数据。

请仅输出相关的表名列表，格式为 JSON 数组，如：["users", "orders"]

---

【回答内容】：
{answer_text}

---

【候选数据库表信息】：
{json.dumps(table_info, ensure_ascii=False, indent=2)}

请分析回答内容是否涉及这些表的字段或业务含义。
"""

    try:
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[
                {"role": "user", "content": prompt}
            ],
            temperature=0.0,
            seed=1,
            max_tokens=200
        )
        content = response.choices[0].message.content.strip()

        # 解析 JSON
        table_names = json.loads(content)
        if isinstance(table_names, list) and all(isinstance(t, str) for t in table_names):
            return table_names
        else:
            print(f"⚠️ 输出格式异常，跳过: {content}")
            return []

    except Exception as e:
        print(f"❌ API 调用失败: {e}")
        return []


# ======================
# 创建关系（不变）
# ======================
def create_uses_table_relationships(answer_id, table_names):
    with driver.session() as session:
        for table_name in table_names:
            session.run("""
                MATCH (a:Answer {id: $answer_id})
                MATCH (t:Table {name: $table_name})
                MERGE (a)-[:USES_TABLE]->(t)
            """, answer_id=answer_id, table_name=table_name)


# ======================
# 主流程
# ======================
def main():
    print("🔍 开始加载 Answer 和 Table 节点...")
    answers = get_all_answers()
    tables = get_all_tables()
    print(f"✅ 加载完成：{len(answers)} 个 Answer，{len(tables)} 个 Table")

    # 预处理 embedding
    for t in tables:
        t['embedding'] = np.array(t['embedding'], dtype=np.float32)

    for i, ans in enumerate(answers):
        print(f"\n🔄 处理 Answer [{i+1}/{len(answers)}]: ID={ans['id']}")

        candidates = find_topk_candidate_tables(ans['embedding'], tables, TOP_K)
        relevant_names = get_relevant_table_names(ans['text'], candidates)
        create_uses_table_relationships(ans['id'], relevant_names)

        print(f"✅ 已关联 {len(relevant_names)} 个表: {relevant_names}")
        time.sleep(0.3)  # 控制请求频率

    print("\n🎉 全部处理完成！")


# ======================
# 启动
# ======================
if __name__ == "__main__":
    try:
        main()
    finally:
        driver.close()